from concurrent.futures import ProcessPoolExecutor, as_completed
import time

from tqdm import tqdm


class MultiprocessingMixin:
    def submit_task(self, executor: ProcessPoolExecutor, task, *args, **kwargs):
        return executor.submit(task, *args, **kwargs)

    def wait_task(self, future_list: list, desc="wait_task"):
        for future in tqdm(
            as_completed(future_list),
            total=len(future_list),
            desc=desc,
        ):
            try:
                future.result()
            except Exception as e:
                print(f"{desc}: {e}")


def worker(x):
    time.sleep(1)  # 模拟耗时任务
    return x * x


if __name__ == "__main__":
    pl = ProcessPoolExecutor(max_workers=4)
    p = MultiprocessingMixin()
    futures = [p.submit_task(pl, worker, i) for i in range(5)]
    p.wait_task(futures, desc="多进程")
